package com.yupi.loveplanneraiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.yupi.loveplanneraiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 抽象基础代理类，用于管理代理状态和执行流程。
 * <p>
 * 提供状态转换、内存管理和基于步骤的执行循环的基础功能。
 * 子类必须实现step方法。
 */
@Data
@Slf4j
public abstract class BaseAgent {

    // 核心属性
    private String name;

    // 提示词
    private String systemPrompt;
    private String nextStepPrompt;

    // 代理状态
    private AgentState state = AgentState.IDLE;

    // 执行步骤控制
    private int currentStep = 0;
    private int maxSteps = 10;

    // LLM 大模型
    private ChatClient chatClient;

    // Memory 记忆（需要自主维护会话上下文）
    private List<Message> messageList = new ArrayList<>();
    
    // 当前会话ID
    private String currentChatId;
    
    // 会话缓存，用于存储不同会话的消息历史
    private static final Map<String, List<Message>> SESSION_CACHE = new ConcurrentHashMap<>();
    
    // 最大缓存消息数（每个会话）
    private static final int MAX_CACHE_MESSAGES = 20;

    /**
     * 运行代理 (同步，实际开发不用)
     *
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    /**
     * 运行代理 (同步，实际开发不用)
     * 
     * @param userPrompt 用户提示词
     * @param chatId 会话ID，用于多轮对话
     * @return 执行结果
     */
    public String run(String userPrompt, String chatId) {
        // 1、基础校验
        if (this.state != AgentState.IDLE) {
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }
        if (StrUtil.isBlank(userPrompt)) {
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }
        
        // 设置当前会话ID
        this.currentChatId = chatId;
        
        // 2、执行，更改状态
        this.state = AgentState.RUNNING;
        
        // 加载会话历史
        loadConversationHistory();
        
        // 记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        
        // 保存结果列表
        List<String> results = new ArrayList<>();
        try {
            // 执行循环
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step {}/{}", stepNumber, maxSteps);
                // 单步执行
                String stepResult = step();
                String result = String.format("🔄 **步骤 %d/%d**\n%s", stepNumber, maxSteps, stepResult);
                results.add(result);
            }
            // 检查是否超出步骤限制
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + maxSteps + ")");
            }
            
            // 保存会话历史
            saveConversationHistory();
            
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("error executing agent", e);
            return "执行错误" + e.getMessage();
        } finally {
            // 3、清理资源
            this.cleanup();
        }
    }
    
    /**
     * 兼容旧接口
     */
    public String run(String userPrompt) {
        return run(userPrompt, null);
    }

    /**
     * 运行代理（流式输出，支持多轮对话）
     *
     * @param userPrompt 用户提示词
     * @param chatId 会话ID，用于多轮对话
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt, String chatId) {
        // 创建一个超时时间较长的 SseEmitter
        SseEmitter sseEmitter = new SseEmitter(300000L); // 5 分钟超时
        // 设置当前会话ID
        this.currentChatId = chatId;
        // 使用线程异步处理，避免阻塞主线程
        CompletableFuture.runAsync(() -> {
            // 1、基础校验
            try {
                if (this.state != AgentState.IDLE) {
                    sseEmitter.send("错误：无法从状态运行代理：" + this.state);
                    sseEmitter.complete();
                    return;
                }
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("错误：不能使用空提示词运行代理");
                    sseEmitter.complete();
                    return;
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
            // 2、执行，更改状态
            this.state = AgentState.RUNNING;
            // 加载会话历史
            loadConversationHistory();
            // 记录新的用户消息
            UserMessage userMessage = new UserMessage(userPrompt);
            messageList.add(userMessage);
            // 保存结果列表
            List<String> results = new ArrayList<>();
            try {
                // 发送初始消息，告知用户AI开始工作
                sseEmitter.send("AI开始思考中...");
                // 执行循环
                for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("Executing step {}/{}", stepNumber, maxSteps);
                    // 发送步骤开始的通知
                    sseEmitter.send("🔄 步骤 " + stepNumber + "/" + maxSteps + "\n");
                    // 单步执行
                    String stepResult = stepWithStream(chunk -> {
                        // 对每个数据块进行SSE格式处理并发送
                        if (chunk != null && !chunk.trim().isEmpty()) {
                            String formattedChunk = chunk + "\n";
                            sseEmitter.send(formattedChunk);
                        }
                    });
                    String result = "步骤 " + stepNumber + ": " + stepResult;
                    results.add(result);
                }
                // 检查是否超出步骤限制
                if (currentStep >= maxSteps) {
                    state = AgentState.FINISHED;
                    String endMessage = "执行结束：达到最大步骤（" + maxSteps + "）";
                    sseEmitter.send(endMessage + "\n");
                }
                // 保存会话历史
                saveConversationHistory();
                // 发送完成信号
                sseEmitter.send("任务DONE\n");
                // 正常完成
                sseEmitter.complete();
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("error executing agent", e);
                try {
                    sseEmitter.send("执行错误：" + e.getMessage() + "\n");
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                // 3、清理资源
                this.cleanup();
            }
        });
        // 设置超时回调
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection timeout");
        });
        // 设置完成回调
        sseEmitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection completed");
        });
        return sseEmitter;
    }
    
    /**
     * 兼容旧接口
     */
    public SseEmitter runStream(String userPrompt) {
        return runStream(userPrompt, null);
    }
    
    /**
     * 执行单个步骤（支持流式输出）
     *
     * @param chunkCallback 数据块回调函数
     * @return 步骤执行结果
     */
    public String stepWithStream(ChunkCallback chunkCallback) {
        // 默认调用普通的step方法
        return step();
    }
    
    /**
     * 用于流式输出的回调接口
     */
    public interface ChunkCallback {
        void onChunk(String chunk) throws Exception;
    }

    /**
     * 定义单个步骤
     *
     * @return
     */
    public abstract String step();

    /**
     * 加载会话历史
     */
    private void loadConversationHistory() {
        if (StrUtil.isNotBlank(currentChatId) && SESSION_CACHE.containsKey(currentChatId)) {
            // 加载缓存的消息历史
            List<Message> cachedMessages = SESSION_CACHE.get(currentChatId);
            this.messageList = new ArrayList<>(cachedMessages);
            log.info("Loaded conversation history for chatId: {}, message count: {}", 
                    currentChatId, cachedMessages.size());
        } else {
            // 如果没有缓存或没有提供chatId，使用空消息列表
            this.messageList = new ArrayList<>();
        }
    }
    
    /**
     * 保存会话历史
     */
    private void saveConversationHistory() {
        if (StrUtil.isNotBlank(currentChatId)) {
            // 限制消息数量，避免内存占用过大
            List<Message> messagesToSave;
            if (messageList.size() > MAX_CACHE_MESSAGES) {
                // 只保留最近的消息
                messagesToSave = messageList.subList(
                        messageList.size() - MAX_CACHE_MESSAGES, 
                        messageList.size()
                );
            } else {
                messagesToSave = messageList;
            }
            // 保存到缓存
            SESSION_CACHE.put(currentChatId, new ArrayList<>(messagesToSave));
            log.info("Saved conversation history for chatId: {}, message count: {}", 
                    currentChatId, messagesToSave.size());
        }
    }
    
    /**
     * 清理资源
     */
    protected void cleanup() {
        // 重置状态
        this.state = AgentState.IDLE;
        this.currentStep = 0;
        // 清空当前实例的消息列表
        this.messageList.clear();
        // 重置当前会话ID
        this.currentChatId = null;
        // 子类可以重写此方法来添加更多清理逻辑
    }

}
